import pickle
import threading
import traceback
from typing import Dict, Optional

import msgpack
import zmq

from ctaf_core.ctaf_objects.base_event import BaseEvent, MsgEvent, reconstruct_event_object_hook
from ctaf_core.ctaf_objects.base_exception import CTAFbaseException, BusCallTimeoutException
from ctaf_core.ctaf_objects.base_service_request import BaseServiceReq
from ctaf_core.ctaf_objects.base_service_response import BaseServiceResp
from ctaf_core.ctaf_utils.config_process_util import load_bus_service_client_settings
from ctaf_core.log_trace_bundle.logging import LogRecorder, LoggerFactory
from ctaf_facade.enumerations.error_code_enum import ErrorCodeEnum
from ctaf_facade.enumerations.service_rsp_code_enum import ResponseCodeEnum


class BusServiceClient:

    log_recorder: LogRecorder = LoggerFactory.get_logger()

    def __init__(self, bus_service_client_setting_filename: str) -> None:
        self.bus_service_client_setting = load_bus_service_client_settings(bus_service_client_setting_filename)

        if not self.bus_service_client_setting:
            BusServiceClient.log_recorder.write_error_log(
                f"BusServiceClient failed to load bus service client settings from {bus_service_client_setting_filename}")
            return

        # rpc_service_name : busServerName
        self.service_in_bus_server: Dict[str, str] = self.get_bus_services_from_setting()

        # busServerName : {"busServerRouterAddress":RouterAddress,"busServerPubAddress":PubAddress}
        self.bus_server_address_info: Dict[str, Dict[str, str]] = self.get_bus_server_address_from_setting()

        # busServerName : ["msgTopic1","msgTopic2" ......]
        self.bus_server_msg_topic: Dict = self.get_bus_server_msg_topic_from_setting()

        # zmq port related
        self._context: zmq.Context = zmq.Context()

        # busServerName : {"socket_dealer":socket_dealer, "socket_sub":socket_sub}
        self.bus_socket_instances: Dict[str, Dict[str, zmq.Socket]] = {}
        self._poller: zmq.Poller = zmq.Poller()  # ZMQ Poller for efficient message reception

        self._active: bool = False  # Client status
        self._event_receiver_thread: Optional[threading.Thread] = None  # event receive thread
        self._rpc_response_thread: Optional[threading.Thread] = None  #  thread for RPC responses

        # --- Concurrency for RPC ---
        self._pending_rpc_requests: Dict[str, threading.Event] = {}  # Event to signal response arrival
        self._rpc_responses: Dict[str, BaseServiceResp] = {}  # Store RPC responses
        self._rpc_lock: threading.Lock = threading.Lock()  # Lock for accessing pending_rpc_requests and rpc_responses

    def get_bus_services_from_setting(self) -> Dict[str, str]:
        service_in_bus_server: Dict[str, str] = {}
        for key, values in self.bus_service_client_setting.items():
            for service in values.get("busServerServices", []):
                service_in_bus_server[service] = key
        return service_in_bus_server

    def get_bus_server_address_from_setting(self) -> Dict[str, Dict[str, str]]:
        bus_server_adr: Dict[str, Dict[str, str]] = {}
        for key, values in self.bus_service_client_setting.items():
            bus_server_adr[key] = {
                "busServerRouterAddress": values.get("busServerRepAddress"),  # Assuming RepAddress is RouterAddress
                "busServerPubAddress": values.get("busServerPubAddress")
            }
        return bus_server_adr

    def get_bus_server_msg_topic_from_setting(self) -> Dict:
        bus_server_msg_topic: Dict = {}
        for key, values in self.bus_service_client_setting.items():
            bus_server_msg_topic[key] = values.get("msgTopic", [])
        return bus_server_msg_topic

    def init_socket_instances(self):
        for key, value in self.bus_server_address_info.items():
            # DEALER socket for requests to ROUTER
            socket_dealer: zmq.Socket = self._context.socket(zmq.DEALER)
            self._set_socket_keep_alive(socket_dealer)

            # SUB socket for subscriptions
            socket_sub: zmq.Socket = self._context.socket(zmq.SUB)
            self._set_socket_keep_alive(socket_sub)

            for topic in self.bus_server_msg_topic.get(key, []):
                socket_sub.setsockopt_string(zmq.SUBSCRIBE, topic)

            # Connect zmq port (DEALER connects to ROUTER)
            router_address = value.get("busServerRouterAddress")
            pub_address = value.get("busServerPubAddress")

            if router_address:
                try:
                    socket_dealer.connect(router_address)
                    self._poller.register(socket_dealer, zmq.POLLIN)  # Register DEALER for polling responses
                except zmq.ZMQError as e:
                    BusServiceClient.log_recorder.write_error_log(
                        f"BusServiceClient failed to connect DEALER socket to {router_address} for {key}: {e}")
                    # Optionally, decide whether to continue or raise
                    continue

            if pub_address:
                try:
                    socket_sub.connect(pub_address)
                    self._poller.register(socket_sub, zmq.POLLIN)  # Register SUB for polling messages
                except zmq.ZMQError as e:
                    BusServiceClient.log_recorder.write_error_log(
                        f"BusServiceClient failed to connect SUB socket to {pub_address} for {key}: {e}")
                    # Optionally, decide whether to continue or raise
                    continue

            server_sockets: Dict[str, zmq.Socket] = {"socket_dealer": socket_dealer, "socket_sub": socket_sub}
            self.bus_socket_instances[key] = server_sockets
            BusServiceClient.log_recorder.write_info_log(f"Bus server: {key} has been connected.")

    def invoke_rpc_request(self, req: BaseServiceReq, timeout: int = 5000) -> BaseServiceResp:
        """
        Invokes a remote procedure call and waits for the response.
        This method is now non-blocking for other RPC calls.
        Args:
            req: The BaseServiceReq object representing the RPC request.
            timeout: The maximum time in milliseconds to wait for a response.
        Returns:
            The BaseServiceResp object.
        Raises:
            RPCRequestTimeout: If the RPC request times out.
            BusClientException: For other RPC related errors.
        """
        called_service_node_name: str = self.service_in_bus_server.get(req.get_called_service_name())
        if not called_service_node_name:
            raise CTAFbaseException(ErrorCodeEnum.NO_TARGET_SERVICE,f"Service '{req.get_called_service_name()}' not found in configuration.")

        socket_dealer: zmq.Socket = self.bus_socket_instances.get(called_service_node_name, {}).get("socket_dealer")
        if not socket_dealer:
            raise CTAFbaseException(ErrorCodeEnum.NO_TARGET_SERVICE,f"DEALER socket not available for bus server: {called_service_node_name}")

        request_id = req.get_ctaf_trace_id()

        # Set up event for this specific request
        response_event = threading.Event()

        rep: BaseServiceResp = BaseServiceResp()
        rep.set_called_service_name(req.get_called_service_name())
        rep.set_ctaf_trace_id(req.get_ctaf_trace_id())

        with self._rpc_lock:
            self._pending_rpc_requests[request_id] = response_event

        try:
            # to serialize RPC request with pickle
            req_bytes = pickle.dumps(req)
            # Send request. DEALER sends a single frame, ROUTER will prepend identity.
            socket_dealer.send_multipart([req_bytes])
            BusServiceClient.log_recorder.write_info_log(
                f"BusServiceClient sent RPC call with trace ID: {request_id} for service: {req.get_called_service_name()}")

            # Wait for the response event with a timeout
            if not response_event.wait(timeout / 1000.0):  # Convert ms to seconds
                raise BusCallTimeoutException(f"BusServiceClient RPC request with trace ID {request_id} timed out after {timeout}ms.")

            # Retrieve the response
            with self._rpc_lock:
                rep = self._rpc_responses.pop(request_id, None)
                if rep is None:
                    raise CTAFbaseException(ErrorCodeEnum.BUS_CALL_ERROR,f"BusServiceClient RPC response for trace ID {request_id} not found after event signal.")
            return rep
        except BusCallTimeoutException:
            BusServiceClient.log_recorder.write_error_log(
                f"BusServiceClient RPC request with trace ID {request_id} to service {req.get_called_service_name()} timed out.")
            rep.set_response_code(ResponseCodeEnum.TIMEOUT_EXECUTE)
            return rep
        except CTAFbaseException as e:
            BusServiceClient.log_recorder.write_error_log(
                f"An CTAFbaseException occurred during RPC call with trace ID {request_id}: {e}\n{traceback.format_exc()}")
            rep.set_response_code(ResponseCodeEnum.FAILED_EXECUTE)
            return rep
        except zmq.ZMQError as e:
            BusServiceClient.log_recorder.write_error_log(
                f"ZMQ error during BusServiceClient RPC call with trace ID {request_id}: {e}\n{traceback.format_exc()}")
            rep.set_response_code(ResponseCodeEnum.FAILED_EXECUTE)
            return rep
        except Exception as e:
            BusServiceClient.log_recorder.write_error_log(
                f"An unexpected error occurred during BusServiceClient RPC call with trace ID {request_id}: {e}\n{traceback.format_exc()}")
            rep.set_response_code(ResponseCodeEnum.FAILED_EXECUTE)
            return rep
        finally:
            # Clean up pending request even if an error occurs
            with self._rpc_lock:
                self._pending_rpc_requests.pop(request_id, None)

    def start(self) -> None:
        """
        Start BusServiceClient
        """
        if self._active:
            BusServiceClient.log_recorder.write_warn_log("BusServiceClient is already active.")
            return

        self.init_socket_instances()

        if not self.bus_socket_instances:
            BusServiceClient.log_recorder.write_error_log(
                "No bus server connections initialized. Cannot start BusServiceClient.")
            return

        self._active = True

        # Start main receive thread (for subscriptions)
        self._event_receiver_thread = threading.Thread(target=self._run_event_receiver, name="BusClientEventReceiver")
        self._event_receiver_thread.start()

        # Start dedicated RPC response thread
        self._rpc_response_thread = threading.Thread(target=self._run_rpc_response_handler,
                                                     name="BusClientRPCResponseHandler")
        self._rpc_response_thread.start()

        BusServiceClient.log_recorder.write_info_log("BusServiceClient started.")

    def _close_socket_instances(self):
        for server_name, socket_item in self.bus_socket_instances.items():
            dealer_socket = socket_item.get("socket_dealer")
            sub_socket = socket_item.get("socket_sub")

            if dealer_socket and not dealer_socket.closed:
                self._poller.unregister(dealer_socket)
                dealer_socket.close()
                BusServiceClient.log_recorder.write_info_log(f"BusServiceClient closed DEALER socket for {server_name}")
            if sub_socket and not sub_socket.closed:
                self._poller.unregister(sub_socket)
                sub_socket.close()
                BusServiceClient.log_recorder.write_info_log(f"BusServiceClient closed SUB socket for {server_name}")
        self.bus_socket_instances.clear()

    def _cleanup_zmq_sockets(self):
        """Cleans up ZMQ sockets and context."""
        self._close_socket_instances()
        if self._context and not self._context.closed:
            self._context.term()
            self._context = None
            BusServiceClient.log_recorder.write_info_log("BusServiceClient ZMQ context terminated.")

    def stop(self) -> None:
        """
        Stop BusServiceClient
        """
        if not self._active:
            BusServiceClient.log_recorder.write_warn_log("BusServiceClient is not active.")
            return

        BusServiceClient.log_recorder.write_info_log("Stopping BusServiceClient...")
        self._active = False  # Signal threads to stop

        self._join_threads()
        self._cleanup_zmq_sockets()

        BusServiceClient.log_recorder.write_info_log("BusServiceClient stopped and threads shut down.")

    def _join_threads(self) -> None:
        """Waits for all client threads to exit."""
        if self._event_receiver_thread and self._event_receiver_thread.is_alive():
            self._event_receiver_thread.join(timeout=1.0)  # Give it some time to finish
            if self._event_receiver_thread.is_alive():
                BusServiceClient.log_recorder.write_warn_log("BusServiceClient event receiver thread did not terminate gracefully.")
        self._thread = None

        if self._rpc_response_thread and self._rpc_response_thread.is_alive():
            self._rpc_response_thread.join(timeout=1.0)  # Give it some time to finish
            if self._rpc_response_thread.is_alive():
                BusServiceClient.log_recorder.write_warn_log(
                    "BusServiceClient RPC response handler thread did not terminate gracefully.")
        self._rpc_response_thread = None

    def _run_event_receiver(self) -> None:
        """
        thread to receive subscription messages using zmq.Poller.
        """
        BusServiceClient.log_recorder.write_info_log("BusClient event receiver thread started.")
        while self._active:
            try:
                # Poll for events with a timeout to allow for graceful shutdown
                socks = dict(self._poller.poll(timeout=100))  # Poll for 100 ms
                for socket_item in self.bus_socket_instances.values():
                    sub_socket = socket_item.get("socket_sub")
                    if sub_socket and sub_socket in socks and socks[sub_socket] == zmq.POLLIN:
                        try:
                            msg_bytes = sub_socket.recv(flags=zmq.NOBLOCK)
                            event: BaseEvent = msgpack.unpackb(
                                msg_bytes, raw=False, object_hook=reconstruct_event_object_hook)
                            if isinstance(event, MsgEvent):
                                #BusServiceClient.log_recorder.write_info_log("A remote event received by BusServiceClient : " + event.get_event_in_str())
                                self.callback(event)
                        except zmq.Again:
                            # No message currently available, continue polling
                            pass
                        except zmq.ContextTerminated:
                            BusServiceClient.log_recorder.write_info_log(
                                "BusClient event receiver ZMQ Context terminated, exiting run loop.")
                            self._active = False
                            break  # Exit inner loop as well
                        except Exception as e:
                            BusServiceClient.log_recorder.write_error_log(
                                f"BusServiceClient error processing subscription message: {e}\n{traceback.format_exc()}")
            except zmq.ContextTerminated:
                BusServiceClient.log_recorder.write_info_log(
                    "BusClient event receiver ZMQ Context terminated, exiting main loop.")
                self._active = False
            except Exception as e:
                BusServiceClient.log_recorder.write_error_log(
                    f"An unexpected error occurred in BusClient event receiver run loop: {e}\n{traceback.format_exc()}")
                self._active = False
        BusServiceClient.log_recorder.write_info_log("BusClient event receiver thread stopped.")

    def _run_rpc_response_handler(self) -> None:
        """
        thread to handle RPC responses from DEALER sockets.
        """
        BusServiceClient.log_recorder.write_info_log("BusClient RPC response handler thread started.")
        while self._active:
            try:
                dealer_sockets = {s["socket_dealer"]: None for s in self.bus_socket_instances.values() if
                                  s.get("socket_dealer")}

                # Create a temporary poller for only DEALER sockets for this loop iteration
                rpc_poller = zmq.Poller()
                for sock in dealer_sockets.keys():
                    rpc_poller.register(sock, zmq.POLLIN)

                socks = dict(rpc_poller.poll(timeout=100))  # Poll for 100 ms

                for socket in dealer_sockets.keys():
                    if socket in socks and socks[socket] == zmq.POLLIN:
                        try:
                            rep_bytes = socket.recv_multipart()
                            rep = pickle.loads(rep_bytes[-1])
                            request_id = None
                            if isinstance(rep, BaseServiceResp):
                                request_id = rep.get_ctaf_trace_id()

                            if request_id:
                                with self._rpc_lock:
                                    self._rpc_responses[request_id] = rep
                                    event = self._pending_rpc_requests.pop(request_id, None)
                                    if event:
                                        event.set()  # Signal that the response has arrived
                                    else:
                                        BusServiceClient.log_recorder.write_warn_log(
                                            f"BusServiceClient received RPC response for unknown or already handled trace ID: {request_id}")
                            else:
                                BusServiceClient.log_recorder.write_warn_log(
                                    f"BusServiceClient received RPC response without a trace ID: {rep}")

                        except zmq.Again:
                            pass  # Should not happen often with poll but good to have
                        except zmq.ContextTerminated:
                            BusServiceClient.log_recorder.write_info_log(
                                "BusClient RPC response handler ZMQ Context terminated, exiting run loop.")
                            self._active = False
                            break  # Exit inner loop as well
                        except Exception as e:
                            BusServiceClient.log_recorder.write_error_log(
                                f"BusServiceClient error processing RPC response: {e}\n{traceback.format_exc()}")
                            # Do not set _active to False here to keep other sockets alive
            except zmq.ContextTerminated:
                BusServiceClient.log_recorder.write_info_log(
                    "BusClient RPC response handler ZMQ Context terminated, exiting main loop.")
                self._active = False
            except Exception as e:
                BusServiceClient.log_recorder.write_error_log(
                    f"An unexpected error occurred in BusClient RPC response handler run loop: {e}\n{traceback.format_exc()}")
                self._active = False
        BusServiceClient.log_recorder.write_info_log("BusClient RPC response handler thread stopped.")

    def callback(self, msg_event: MsgEvent) -> None:
        """
        Placeholder for application-specific event handling.
        This method should be overridden by the user.
        """
        # print(f"Received event: {msg_event.get_event_in_str()}") # For debugging purposes
        pass

    @staticmethod
    def _set_socket_keep_alive(socket_instance: zmq.Socket):
        """Helper to set TCP keep-alive options on a socket."""
        try:
            socket_instance.setsockopt(zmq.TCP_KEEPALIVE, 1)
            socket_instance.setsockopt(zmq.TCP_KEEPALIVE_IDLE,
                                       60)  # Idle time before sending first keepalive probe (seconds)
            socket_instance.setsockopt(zmq.TCP_KEEPALIVE_INTVL, 30)  # Interval between keepalive probes (seconds)
            socket_instance.setsockopt(zmq.TCP_KEEPALIVE_CNT,
                                       5)  # Number of keepalive probes before declaring connection dead
        except zmq.ZMQError as e:
            BusServiceClient.log_recorder.write_warn_log(f"BusServiceClient failed to set TCP keep-alive options: {e}")

    @staticmethod
    def sub_msg_topic(socket_instance: zmq.Socket, topic: str) -> None:
        """Helper to subscribe to a message topic."""
        socket_instance.setsockopt_string(zmq.SUBSCRIBE, topic)